{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import and namelist\n",
    "import xarray as xr\n",
    "import sys; sys.path.append('/work/home/linjin/met/work')\n",
    "import numpy as np\n",
    "# namelist\n",
    "path = '/work/home/linjin/met/met/grid/plev_linjin_2025062200_000_cf.nc'\n",
    "heights = [\n",
    "    250, 500, 750, 1000, 1250, 1500, 1750, 2000, 2250, 2500,\n",
    "    2750, 3000, 3250, 3500, 3750, 4000, 4250, 4500, 4750, 5000,\n",
    "    5250, 5500, 5750, 6000, 6250, 6500, 6750, 7000, 7250, 7500,\n",
    "    7750, 8000, 8250, 8500, 8750, 9000, 9250, 9500, 9750, 10000,\n",
    "    10250, 10500, 10750, 11000, 11250, 11500, 11750, 12000, 12250, 12500,\n",
    "    12750, 13000, 13250, 13500, 13750, 14000, 14250, 14500, 14750, 15000,\n",
    "    15250, 15500, 15750, 16000, 16250, 16500, 16750, 17000, 17250, 17500,\n",
    "    17750, 18000, 18250, 18500, 18750, 19000, 19250, 19500, 19750, 20000,\n",
    "    20500, 21000, 21500, 22000, 22500, 23000, 23500, 24000, 24500, 25000,\n",
    "    25500, 26000, 26500, 27000, 27500, 28000, 28500, 29000, 29500, 30000,\n",
    "    30500, 31000, 31500, 32000, 32500, 33000, 33500, 34000, 34500, 35000,\n",
    "    35500, 36000, 36500, 37000, 37500, 38000, 38500, 39000, 39500, 40000,\n",
    "    40500, 41000, 41500, 42000, 42500, 43000, 43500, 44000, 44500, 45000,\n",
    "    45500, 46000, 46500, 47000, 47500, 48000, 48500, 49000, 49500, 50000,\n",
    "    50500, 51000, 51500, 52000, 52500, 53000, 53500, 54000, 54500, 55000,\n",
    "    55500, 56000, 56500, 57000, 57500, 58000, 58500, 59000, 59500, 60000,\n",
    "    60500, 61000, 61500, 62000, 62500, 63000, 63500, 64000, 64500, 65000,\n",
    "    65500, 66000, 66500, 67000, 67500, 68000, 68500, 69000, 69500, 70000\n",
    "]\n",
    "f = xr.open_dataset(path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f.valid_time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gsj_plot import see\n",
    "path = \"/work/home/linjin/combine_run/COMBINE_WORK_ROOT/2025062200/hlev/hlev_linjin_2025062200_024.nc\"\n",
    "f = xr.open_dataset(path)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u = f.u.values[0]     # m/s\n",
    "v = f.v.values[0]     # m/s\n",
    "WindSpeed = np.sqrt(u**2+v**2)\n",
    "h = f.z.values[0]/9.8    # m\n",
    "t = f.t.values[0]     # K\n",
    "p = f.pressure_level.values   # 33 level   (33,)  hPa\n",
    "rho = p[:,None,None]*100/287/t     #  (33,lat,lon)  kg/m^3\n",
    "latitude = f.latitude.values\n",
    "longitude = f.longitude.values\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.interpolate import CubicSpline\n",
    "import numpy as np\n",
    "from multiprocessing import Pool\n",
    "\n",
    "\n",
    "def interp_column(j,i,z_coord,col):\n",
    "    \"\"\"\n",
    "    对单条垂直廓线做 cubic 插值\n",
    "    Inputs:\n",
    "        z_coords是单条廓线的垂直高度数组 e.g. (33,)\n",
    "        col 是某变量在该廓线上的值       e.g.  (33,)\n",
    "        heights:   需要插值的目标高度    e.g.   (100,)\n",
    "        j : index 用来恢复数组\n",
    "        i : index 用来恢复数组\n",
    "    return: \n",
    "        目标高度值           e.g. (100,)\n",
    "    \"\"\"\n",
    "    cs = CubicSpline(z_coord, col)\n",
    "    return j,i,cs(heights)                    \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "t_interp    = np.zeros((len(heights),len(latitude),len(longitude)))\n",
    "rho_interp  = np.zeros((len(heights),len(latitude),len(longitude)))\n",
    "WindSpeed_interp = np.zeros((len(heights),len(latitude),len(longitude)))\n",
    "nlat,nlon = latitude.size, longitude.size\n",
    "\n",
    "# T\n",
    "inputs = [(j,i,h[:,j,i],t[:,j,i]) for j in range(nlat) for i in range(nlon)]\n",
    "with Pool() as pool:\n",
    "    results = pool.map(interp_column, inputs)\n",
    "for result in results:\n",
    "    j, i, interp_col = result\n",
    "    t_interp[:, j, i] = interp_col\n",
    "\n",
    "# rho\n",
    "inputs = [(j,i,h[:,j,i],rho[:,j,i]) for j in range(nlat) for i in range(nlon)]\n",
    "with Pool() as pool:\n",
    "    results = pool.map(interp_column, inputs)\n",
    "for result in results:\n",
    "    j, i, interp_col = result\n",
    "    rho_interp[:, j, i] = interp_col\n",
    "\n",
    "# WindSpeed\n",
    "inputs = [(j,i,h[:,j,i],WindSpeed[:,j,i]) for j in range(nlat) for i in range(nlon)]\n",
    "with Pool() as pool:\n",
    "    results = pool.map(interp_column, inputs)\n",
    "for result in results:\n",
    "    j, i, interp_col = result\n",
    "    WindSpeed_interp[:, j, i] = interp_col\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# np2dataset\n",
    "from datetime import datetime\n",
    "\n",
    "nt, nz, nlat, nlon = 1, len(heights), latitude, longitude\n",
    "\n",
    "ds_out = xr.Dataset(\n",
    "    {\n",
    "        'temperature': (['time', 'height', 'lat', 'lon'], t_interp[np.newaxis, ...],\n",
    "                        {'units': 'K',\n",
    "                         'long_name': 'temperature',\n",
    "                         'coordinates': 'lat lon',\n",
    "                         'grid_mapping': 'grid_mapping',\n",
    "                         '_FillValue': -9999.}),\n",
    "        'density': (['time', 'height', 'lat', 'lon'], rho_interp[np.newaxis, ...],\n",
    "                    {'units': 'kg m-3',\n",
    "                     'long_name': 'air density',\n",
    "                     'coordinates': 'lat lon',\n",
    "                     'grid_mapping': 'grid_mapping',\n",
    "                     '_FillValue': -9999.}),\n",
    "        'wind_speed': (['time', 'height', 'lat', 'lon'], WindSpeed_interp[np.newaxis, ...],\n",
    "                       {'units': 'm s-1',\n",
    "                        'long_name': 'wind speed',\n",
    "                        'coordinates': 'lat lon',\n",
    "                        'grid_mapping': 'grid_mapping',\n",
    "                        '_FillValue': -9999.}),\n",
    "    },\n",
    "    coords={\n",
    "        'time': (['time'], [f.time.values],  # 可替换为真实 init/valid time\n",
    "                 {'standard_name': 'time',\n",
    "                  'units': 'seconds since 1970-01-01 00:00:00'}),\n",
    "        'height': (['height'], heights,\n",
    "                   {'standard_name': 'height',\n",
    "                    'units': 'm',\n",
    "                    'positive': 'up'}),\n",
    "        'lat': (['lat'], h.latitude.values,\n",
    "                {'standard_name': 'latitude',\n",
    "                 'units': 'degrees_north'}),\n",
    "        'lon': (['lon'], h.longitude.values,\n",
    "                {'standard_name': 'longitude',\n",
    "                 'units': 'degrees_east'}),\n",
    "        'grid_mapping': ([], 0,  # 标量变量\n",
    "                         {'grid_mapping_name': 'latitude_longitude',\n",
    "                          'semi_major_axis': 6371000.,\n",
    "                          'inverse_flattening': 0.})\n",
    "    },\n",
    "    attrs={\n",
    "        'Conventions': 'CF-1.8',\n",
    "        'title': 'FNL vertical interpolation to height levels',\n",
    "        'institution': 'For zkxt Full Atmsphere verification',\n",
    "        'history': f'{datetime.utcnow().isoformat()} created',\n",
    "    }\n",
    ")\n",
    "\n",
    "# ------------------------------------------------------------------\n",
    "# 2. 写出 NetCDF\n",
    "# ------------------------------------------------------------------\n",
    "out_file = '/work/home/linjin/met/work/grid/ERA5_20250623_0000_pl_height.nc'\n",
    "ds_out.to_netcdf(out_file, format='NETCDF4')\n",
    "print('CF-compliant NetCDF saved to:', out_file)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "met_3.10",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
